ZeRO Stage 3 ব্যবহার করে মডেল অপ্টিমাইজ করা

ZeRO Stage 3 (Zero Redundancy Optimizer) DeepSpeed এর একটি শক্তিশালী বৈশিষ্ট্য যা মডেল অপ্টিমাইজেশনে গুরুত্বপূর্ণ ভূমিকা পালন করে। এটি বৃহৎ মডেল প্রশিক্ষণে GPU মেমোরি ব্যবহার উল্লেখযোগ্যভাবে কমিয়ে আনে এবং প্রশিক্ষণের জন্য কার্যকরী সমাধান প্রদান করে। ZeRO Stage 3 এর মাধ্যমে মডেলের parameter states, gradients, এবং optimizer states সবকিছুই বিভিন্ন GPU তে শার্ড করা হয়, যা মেমোরি ব্যবহারে সর্বাধিক দক্ষতা আনে।

ZeRO Stage 3 ব্যবহার করে মডেল অপ্টিমাইজ করার ধাপসমূহ

নিচে ZeRO Stage 3 ব্যবহার করে একটি মডেল অপ্টিমাইজ করার সম্পূর্ণ প্রক্রিয়া দেখানো হলো।

ধাপ ১: প্রয়োজনীয় লাইব্রেরি ইনস্টল করা

প্রথমে আপনাকে নিশ্চিত করতে হবে যে আপনার সিস্টেমে DeepSpeed এবং অন্যান্য প্রয়োজনীয় লাইব্রেরিগুলি ইনস্টল করা আছে।

pip install torch torchvision deepspeed

ধাপ ২: মডেল তৈরি করা

এখন একটি কাস্টম নিউরাল নেটওয়ার্ক মডেল তৈরি করুন। উদাহরণস্বরূপ, আমরা একটি সাধারণ Fully Connected Neural Network ব্যবহার করব।

import torch
import torch.nn as nn
import deepspeed

class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(10, 64)
        self.fc2 = nn.Linear(64, 1)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

# মডেল ইনিশিয়ালাইজেশন
model = SimpleNN()

ধাপ ৩: DeepSpeed কনফিগারেশন তৈরি করা

ZeRO Stage 3 এর জন্য একটি কনফিগারেশন ফাইল তৈরি করুন। এই কনফিগারেশনটি মডেলের মেমোরি ব্যবহারের জন্য প্রয়োজনীয় সেটিংস ধারণ করবে।

{
    "train_batch_size": 32,
    "fp16": {
        "enabled": true  # Mixed Precision Training
    },
    "zero_optimization": {
        "stage": 3  # ZeRO Stage 3
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": 5e-5
        }
    }
}

ধাপ ৪: ডেটাসেট তৈরি করা

মডেলের প্রশিক্ষণের জন্য ডেটাসেট তৈরি করুন। এখানে একটি সিম্পল সিমুলেটেড ডেটাসেট ব্যবহার করা হবে।

# সিমুলেটেড ডেটাসেট তৈরি
X = torch.randn(1000, 10)  # 1000 নমুনা, 10 বৈশিষ্ট্য
y = torch.randn(1000, 1)   # টার্গেট আউটপুট

# DataLoader তৈরি করা
train_loader = torch.utils.data.DataLoader(list(zip(X, y)), batch_size=32, shuffle=True)

ধাপ ৫: DeepSpeed ইনিশিয়ালাইজেশন

DeepSpeed এর মাধ্যমে মডেল এবং অপটিমাইজার ইনিশিয়ালাইজ করুন।

# DeepSpeed ইনিশিয়ালাইজেশন
deepspeed_config = "deepspeed_config.json"

model_engine, optimizer, _, _ = deepspeed.initialize(
    config=deepspeed_config,
    model=model,
    model_parameters=model.parameters()
)

ধাপ ৬: কাস্টম ট্রেইনিং লুপ তৈরি করা

মডেল প্রশিক্ষণের জন্য একটি কাস্টম ট্রেইনিং লুপ তৈরি করুন।

# Training Loop
num_epochs = 5
for epoch in range(num_epochs):
    model_engine.train()
    running_loss = 0.0

    for inputs, labels in train_loader:
        optimizer.zero_grad()  # গ্র্যাডিয়েন্টস রিসেট করা

        # Forward pass
        outputs = model_engine(inputs)
        loss = nn.MSELoss()(outputs, labels)

        # Backward pass
        model_engine.backward(loss)
        model_engine.step()  # অপটিমাইজার স্টেপ

        running_loss += loss.item()

    epoch_loss = running_loss / len(train_loader)
    print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {epoch_loss:.4f}')

ধাপ ৭: মডেল সেভ এবং লোড করা

মডেল এবং অপটিমাইজার স্টেট সেভ এবং লোড করার জন্য save_checkpoint এবং load_checkpoint ফাংশন ব্যবহার করুন।

# মডেল সেভ করা
model_engine.save_checkpoint("./checkpoints")

# মডেল লোড করা
model_engine.load_checkpoint("./checkpoints")

উপসংহার

ZeRO Stage 3 ব্যবহার করে মডেল অপ্টিমাইজ করা মডেলের GPU মেমোরি ব্যবহারে উল্লেখযোগ্য সুবিধা এনে দেয় এবং বৃহৎ মডেল প্রশিক্ষণের জন্য কার্যকরী সমাধান প্রদান করে। এই প্রক্রিয়ার মাধ্যমে, আপনি মডেল প্রশিক্ষণ করতে পারেন, মডেল সেভ এবং লোড করতে পারেন, এবং DeepSpeed এর মাধ্যমে পারফরম্যান্স বিশ্লেষণ করতে পারেন।